#!/usr/bin/python

"""

Concept:
Set = 13655
A = {}
B = {}
A intersect B = {}
A union B = {}

Usage: dissertation_CountSharedOrthologs.py CYANO.orthologs.txt GKIL 58011
Note: Regex is hard-coded for now.

Author: Jimmy Saw
Date: 05-03-2012

"""

import sys
import re

def unique(a):
    """ return the list with duplicate elements removed """
    return list(set(a))

def intersect(a, b):
    """ return the intersection of two lists """
    return list(set(a) & set(b))

def union(a, b):
    """ return the union of two lists """
    return list(set(a) | set(b))

orthofile = open(sys.argv[1], "rU")
orthos = orthofile.readlines()

organism1 = str(sys.argv[2])

organism2 = str(sys.argv[3])

org1 = re.compile('GKIL\|\w+')
org2 = re.compile('58011\|\w+')
orgs = re.compile('(\w+)\|\w+')

totalorths = []
organismlist = []
o1hits = []
o2hits = []
o1count = 0
o2count = 0

olist = []
for line in orthos:
    c = line.split(' ')
    oid = c[0]
    ogroups = c[1:]
    totalorths.append((oid, 0))
    olist.append((oid, ogroups))

for j in olist:
    oid = j[0]
    tmplist = j[1]
    for k in tmplist:
        if orgs.match(k):
            p = orgs.match(k)
            organismlist.append(p.group(1))

uniqueorgs = len(list(set(organismlist)))
uniqueorglist = list(set(organismlist))

orgsdict = {}
ocountsdict = {}
orgsorthdict = {}
for orgname in uniqueorglist:
    orgsdict[orgname] = []
    ocountsdict[orgname] = 0
    orgsorthdict[orgname] = []

for i in uniqueorglist:
    for j in olist:
        oid = j[0]
        tmplist = j[1]
        for k in tmplist:
            ocount = 0
            tmpstr = k.split('|')
            if i == tmpstr[0]:
                ocount = 1
                orgsdict[i].append((oid, ocount))
                ocountsdict[i] += 1
                orgsorthdict[i].append(oid)

testintersect1 = intersect(orgsorthdict[organism1], orgsorthdict[organism2])
testunion1 = union(orgsorthdict[organism1], orgsorthdict[organism2])
print "Total number of orthologs:", len(totalorths)
print organism1, "has", len(orgsorthdict[organism1]), "orthologs out of", len(totalorths), "orthologs."
print organism2, "has", len(orgsorthdict[organism2]), "orthologs out of", len(totalorths), "orthologs."
print organism1 + " vs " + organism2, "Intersection:", len(testintersect1)
print organism1 + " vs " + organism2, "Union:", len(testunion1)

"""
#print phyletic patterns here
for i in uniqueorglist:
    toprintlist = []
    for j in totalorths:
        if j[0] in orgsorthdict[i]:
            #totalorths[i] = ((j[0], 1))
            toprintlist.append(str(1))
        else:
            toprintlist.append(str(0))
    print '>' + i
    print ''.join(toprintlist)
"""

orthofile.close()
